from GNN.dataset import  GraphDataset
from GNN.function import Trainer
from GNN.models import GNNModel

import torch,os,json,argparse
from torch.utils.data import DataLoader

def get_configs():
    '''
    参数
    data_path: 数据集路径
    max_epoch: 最大训练轮数
    
    '''
    parser = argparse.ArgumentParser(description="Configuration for ROP dataset processing and model training")
    parser.add_argument('--data_path', type=str, default='../Dataset/infantImages', help="Path to the dataset")
    parser.add_argument('--max_epoch', type=int, default=10, help="Number of epochs")
    
    args = parser.parse_args()
    return args
def main():
    args = get_configs()
    data_path = args.data_path
    max_epoch = args.max_epoch
    
    model = GNNModel()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    loss_fn = torch.nn.MSELoss()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    dataset = GraphDataset(os.path.join(data_path, 'graph_record.json'),os.path.join(data_path, 'split', 'all.json'),'train')
    data_loader = DataLoader(dataset, batch_size=1, shuffle=True)
    trainer = Trainer(model, optimizer, loss_fn, data_loader, device)
    
    trainer.train(max_epoch)